# 定义 admin_required 装饰器  
from functools import wraps

from flask import jsonify, redirect, session, url_for

from app.doctor.models import Doctor


def doctor_required(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        # authentication
        if 'doctor_id' not in session:
            return jsonify({"code":1, 'message': 'doctor_id not in session'}), 401

        doctor = Doctor.query.get(session['doctor_id'])
        if not doctor:
            return jsonify({"code":1, 'message': 'Doctor not found'}), 401

        # 如果验证通过，则执行原始函数
        return func(*args, **kwargs)
    
    return wrapper


def doctor_required_web(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if 'doctor_id' in session:
            if Doctor.query.get(session['doctor_id']):
                return func(*args, **kwargs)

        return redirect(url_for('doctor.login_doctor_web'))

    return wrapper
